"""
数据采集模块
用于从各种数据源采集长尾关键词
"""
import pandas as pd
import logging
from typing import List, Dict, Set
import time
from .utils import random_delay
from .data_sources import DataSourceManager


class KeywordCollector:
    """关键词采集器"""
    
    def __init__(self, config: Dict):
        """
        初始化采集器
        
        Args:
            config: 配置字典
        """
        self.config = config
        self.data_config = config.get('data_collection', {})
        self.request_config = self.data_config.get('request_config', {})
        self.logger = logging.getLogger(__name__)
        self.collected_keywords = set()
        
        # 初始化数据源管理器
        self.source_manager = DataSourceManager(config)
    
    def expand_keywords(self, seed_keywords: List[str]) -> Set[str]:
        """
        扩展种子关键词
        
        Args:
            seed_keywords: 种子关键词列表
        
        Returns:
            扩展后的关键词集合
        """
        all_keywords = set(seed_keywords)
        
        for seed_keyword in seed_keywords:
            self.logger.info(f"正在处理种子关键词: {seed_keyword}")
            
            # 从所有启用的数据源获取关键词
            keywords = self.source_manager.collect_from_all_sources(seed_keyword)
            all_keywords.update(keywords)
            
            # 添加延迟
            delay_range = self.request_config.get('delay_range', [1, 3])
            random_delay(delay_range)
        
        return all_keywords
    
    def collect_keywords(self) -> pd.DataFrame:
        """
        执行关键词采集
        
        Returns:
            包含采集结果的DataFrame
        """
        self.logger.info("开始采集长尾关键词...")
        
        # 获取种子关键词
        seed_keywords = self.data_config.get('seed_keywords', [])
        if not seed_keywords:
            raise ValueError("未配置种子关键词")
        
        # 扩展关键词
        expanded_keywords = self.expand_keywords(seed_keywords)
        
        # 转换为DataFrame
        keywords_data = []
        for i, keyword in enumerate(expanded_keywords, 1):
            keywords_data.append({
                'id': i,
                'keyword': keyword,
                'length': len(keyword),
                'source': 'collected',
                'collected_time': pd.Timestamp.now()
            })
        
        df = pd.DataFrame(keywords_data)
        
        self.logger.info(f"关键词采集完成，共采集到 {len(df)} 个关键词")
        
        return df
    
    def get_source_info(self) -> List[Dict]:
        """获取所有数据源信息"""
        return self.source_manager.get_source_info()
    
    def add_data_source(self, source):
        """添加新的数据源"""
        self.source_manager.add_source(source)
    
    def remove_data_source(self, source_name: str):
        """移除数据源"""
        self.source_manager.remove_source(source_name)


def main():
    """测试函数"""
    from utils import load_config, setup_logging
    
    # 加载配置
    config = load_config()
    logger = setup_logging(config)
    
    # 创建采集器
    collector = KeywordCollector(config)
    
    # 显示数据源信息
    source_info = collector.get_source_info()
    print("当前可用的数据源:")
    for source in source_info:
        print(f"  - {source['name']} ({source['type']}) - 启用: {source['enabled']}")
    
    # 采集关键词
    keywords_df = collector.collect_keywords()
    
    # 保存结果
    keywords_df.to_csv('data/raw/collected_keywords.csv', 
                      index=False, encoding='utf-8-sig')
    
    print(f"采集完成，共获得 {len(keywords_df)} 个关键词")
    print(keywords_df.head())


if __name__ == "__main__":
    main()
